#-*- coding:utf-8 -*-
import tensorflow as tf
from getStrokeLabel import genStrokes
#from getStroke import genStrokes
import numpy as np
from config import Config
from myCNN import CNN2
import os
import h5py
import matplotlib.pyplot as plt
import cv2


def read_data(path):
     images = np.array(h5py.File(path, 'r').get('images'))
     labels = np.array(h5py.File(path, 'r').get('labels'))
     return images, labels

def read_matte(path):
     mattes = np.array(h5py.File(path, 'r').get('matte'))
     return mattes


config = Config()
n_steps = config.step
draw = genStrokes(config, )
cnn2 = CNN2(config,)


def MSE(image,gt):
      if len(image.shape)>2:
           image=image[:,:,0]
      if len(gt.shape)>2:
           gt=gt[:,:,0]
  
      width=image.shape[0]
      height=image.shape[1]
      sumup=0
      for j in range(width):
         for k in range(height):
             sumup=sumup+(image[j,k]-gt[j,k])*(image[j,k]-gt[j,k])
      MSE=sumup/(width*height)
      return MSE
    

matte_dir = os.path.join(os.getcwd(), 'matte.h5')
data_dir = os.path.join(os.getcwd(), 'train.h5')
train_images, train_alphas = read_data(data_dir)
#matte_arr=read_matte(matte_dir)
strokeMaps = train_images.copy()
images=train_images.copy()

'''
this_matte=matte_arr[0,0,:,:]
gt=train_alphas[0]
this_mse=MSE(this_matte,gt)
#cv2.imshow('matte',np.uint8(this_matte))
#cv2.waitKey(0)
#cv2.destroyAllWindows()

for i in range (config.num_glimpses):
     mattes=matte_arr[i]
     matte=mattes[0]
     #cv2.imshow('matte',np.uint8(matte))
     #cv2.waitKey(0)
     #cv2.destroyAllWindows()
     if MSE(matte,gt)!=this_mse:
         this_mse=MSE(matte,gt)
         print i , this_mse
'''


flag=True
counter=0
RandomMse=[]
while(counter<75):
    loc = np.random.uniform(-1.0,1.0,(train_alphas.shape[0],2))
    strokeMaps = draw.drawstroke(strokeMaps,loc,train_alphas)
    alpha_result = cnn2.op(train_images,strokeMaps)
    Randomalpha=alpha_result[0]
    counter=counter+1
    #RandomMse.append(MSE(Randomalpha,gt))

print '########################'
#print RandomMse


#for i in range (config.num_glimpses):
    #loc = np.random.uniform(-1.0,1.0,(train_alphas.shape[0],2))
    #strokeMaps = draw.drawstroke(strokeMaps,loc,train_alphas)





#alpha_result = cnn2.op(train_images,strokeMaps)

  
